#include "common.h"
#include "kernel_operator.h"
#ifdef __CCE_KT_TEST__
#include <iostream>
#endif
namespace AscendC
{
template <typename _DT_PREDICT, typename _DT_LABEL, typename _DT_DOUT, typename _DT_Y> class Kernel
{
  public:
    TPipe pipe;
    CommonTiling tilingData;

    bool isMean = true;
    float cof = 2;
    bool dout_sameshape = false;
    _DT_DOUT dout_value;

    DefInTensor(PREDICT);
    DefInTensor(LABEL);
    DefInTensor(DOUT);
    DefOutTensor(Y);

  public:
    __aicore__ inline Kernel()
    {
    }
    __aicore__ inline void Init(GM_ADDR predict, GM_ADDR label, GM_ADDR dout, GM_ADDR y, bool is_mean,
                                bool doutSameShape, CommonTiling &tiling)
    {
        tilingData = tiling;
        this->isMean = is_mean;
        if (is_mean)
        {
            cof /= float(tilingData.size);
        }
        dout_sameshape = doutSameShape;
        // 切分global
        SetGlobalBuffer(PREDICT, predict, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(LABEL, label, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(DOUT, dout, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(PREDICT, tilingData.bufferSize);
        InitQueueSimple(LABEL, tilingData.bufferSize);
        InitQueueSimple(DOUT, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);

        if (!dout_sameshape)
        {
            auto dout_ptr = CAST_GM_ADDR(TypeOf(DOUT), dout, 0);
            dout_value = dout_ptr[0];
        }
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(i, tilingData.bufferSize);
            CopyOut(i, tilingData.bufferSize);
        }
        auto copy_count = tilingData.getFinalCopyCount<TypeOf(PREDICT)>();
        CopyIn(finnal_progress, copy_count);
        Compute(finnal_progress, tilingData.finalCalcCount);
        CopyOut(finnal_progress, copy_count);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(PREDICT, index, calcCount);
        EnQueGlobal2Local(LABEL, index, calcCount);
        EnQueGlobal2Local(DOUT, index, calcCount);
    }
    /*
    if 'mean' == reduction:
        reduce_elts = 1.0
        for i in input_predict.shape:
            reduce_elts *= i
        cof = (reduce_elts**(-1)) * 2.0
    else:
        cof = 2.0
    sub_res = input_predict - input_label
    norm_grad = sub_res * cof
    golden = norm_grad * input_dout
    */
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(PREDICT);
        DeQueSimple(LABEL);
        DeQueSimple(DOUT);
        QueAllocSimple(Y);
        // sub_res = input_predict - input_label
        Sub(LTensorName(Y), LTensorName(PREDICT), LTensorName(LABEL), calcCount);
        // norm_grad = sub_res * cof
        Muls(LTensorName(Y), LTensorName(Y), _DT_PREDICT(cof), calcCount);
        // golden = norm_grad * input_dout
        if (dout_sameshape)
        {
            Mul(LTensorName(Y), LTensorName(Y), LTensorName(DOUT), calcCount);
        }
        else
        {
            Muls(LTensorName(Y), LTensorName(Y), dout_value, calcCount);
        }
        EnQue(Y);
        QueFree(DOUT);
        QueFree(LABEL);
        QueFree(PREDICT);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount);
    }
};

} // namespace AscendC
extern "C" __global__ __aicore__ void mse_loss_grad(GM_ADDR predict, GM_ADDR label, GM_ADDR dout, GM_ADDR y,
                                                    GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    AscendC::Kernel<DTYPE_PREDICT, DTYPE_LABEL, DTYPE_DOUT, DTYPE_Y> op;
    op.Init(predict, label, dout, y, tiling_data.tag == 0, tiling_data.doutSameShape, commonTiling);
    op.Process();
}